# 导包.
import numpy as np
import pandas as pd
from sklearn.tree import DecisionTreeRegressor, ExtraTreeRegressor  # 回归决策树
from sklearn.linear_model import LinearRegression  # 线性回归
import matplotlib.pyplot as plt  # 绘图
plt.rcParams['font.sans-serif'] = ['SimHei'] # 正常显示汉字
plt.rcParams['axes.unicode_minus'] = False   # 正常显示符号
# 1. 获取数据.
x = np.array(list(range(1, 11))).reshape(-1, 1)
y = np.array([5.56, 5.70, 5.91, 6.40, 6.80, 7.05, 8.90, 8.70, 9.00, 9.05])

# 2. 创建线性回归 和 决策树回归模型.
m1 = LinearRegression()
m2 = DecisionTreeRegressor(max_depth=1)
m3 = DecisionTreeRegressor(max_depth=3)

# 3. 训练模型.
m1.fit(x,y)
m2.fit(x,y)
m3.fit(x,y)

# 4. 准备测试数据, 用于测试.
# 起始, 结束, 步长.
x_test = np.arange(0.0, 10.0, 0.1).reshape(-1, 1)

# 5. 模型预测.
m1.predict(x_test)
m2.predict(x_test)
m3.predict(x_test)
# 6. 绘图 plt.figure(figsize=(x,y))
plt.figure(figsize=(20,20))
# 散点图(原始的坐标) : .scatter
plt.scatter(x, y)
# 线性回归的预测结果: .plot()
plt.plot(x_test, m1.predict(x_test),color='red')
# 回归决策树  .plot()
plt.plot(x_test, m2.predict(x_test),color='black')
plt.plot(x_test, m3.predict(x_test),color='green')

# 显示图例.
plt.legend('图')
# 设置x轴标签: .xlabel()
plt.xlabel('x')
# 设置y轴标签: .ylabel()
plt.ylabel('y')
# 设置标题
plt.title(label='线性回归与决策树')

# 显示图片
plt.show()